/**
 * Copyright 2012 Alexey Ragozin
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.gridkit.benchmark.gc;

import java.io.IOException;
import java.io.Serializable;
import java.lang.management.GarbageCollectorMXBean;
import java.lang.management.ManagementFactory;
import java.lang.management.MemoryPoolMXBean;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.TimeUnit;

import com.beust.jcommander.Parameter;
import com.beust.jcommander.Parameters;

/**
 * @author Alexey Ragozin (alexey.ragozin@gmail.com)
 */
@SuppressWarnings("restriction")
@Parameters(commandDescriptionKey="ygc", commandDescription = "Young GC benchmark")
public class YoungGCPauseBenchmark {

	public enum DataMode {
		STRING,
		LONG,
		INT,
		CONST,
	}
	
	private Random random = new Random(0);
	private Map<Integer, Object> maps[];

	@Parameter(names = {"-es", "--entry-size"}, description = "Memory size of entry")
	int entrySize = -1;
	
	@Parameter(names = {"-h", "--head-room"}, description = "Reserved portion of old space in MiB")
	int headRoom = 256;
	
	@Parameter(names = {"-m", "--data-mode"}, description = "Type of garbage data")
	DataMode mode = DataMode.STRING;

	@Parameter(names = {"-l", "--string-len"}, description = "Average length of string")
	int stringLen = 64;
	
	@Parameter(names = {"-d", "--dry-mode"}, description = "After filling old space, test will stop dirting old space. This mode should exclude dirty page scanning from measurement.")
	boolean dryMode = false;

	@Parameter(names = {"--max-time"}, description = "Benchmark time limit (sec)")
	int maxTime = 60;
	
	@Parameter(names = {"--max-old"}, description = "Number of old collections to benchmark")
	int maxOld = -1;

	@Parameter(names = {"--max-young"}, description = "Number of young collections to benchmark")
	int maxYoung = -1;
	
	@Parameter(names = {"--min-time"}, description = "Minimum benchmark time limit (sec)")
	int minTime = -1;
	
	@Parameter(names = {"--min-old"}, description = "Minimum number of old collections to measure")
	int minOld = -1;
	
	@Parameter(names = {"--min-young"}, description = "Minimum number of young collections to measure")
	int minYoung = -1;
	
	@Parameter(names = {"-r", "--overide-rate"}, description = "Chance that new object would be put to the map (otherwise existing would be reiserted)")
	double overrideRate = 0.1;
	
	@Parameter(names = {"-e", "--print-events"}, description = "Print GC events to console")
	boolean printEvents = false;

	private int count;
	
	private com.sun.management.GarbageCollectorMXBean oldGcMBean;
	private GarbageCollectorMXBean youngGcMBean;
	
	private MemoryPoolMXBean oldMemPool;
	private double activeOverrideRate = 1.1;
	
	
	private static char[] STRING_TEMPLATE;
	
	private static List<String> OLD_POOLS = Arrays.asList("Tenured Gen", "PS Old Gen", "CMS Old Gen", "G1 Old Gen", "Old Space"); 
	private static List<String> CONC_MODE = Arrays.asList("CMS Old Gen", "G1 Old Gen");
	private boolean concurentMode;
	
	private int align(int size, int al) {
		return (size + al - 1) & (~(al - 1));
	}
	
	public TestResult benchmark() throws IOException {
		
		STRING_TEMPLATE = new char[stringLen];
		
		if (mode == DataMode.CONST) {
			overrideRate = 1.1d;
		}

		System.out.println("Java: " + System.getProperty("java.version") + " VM: " + System.getProperty("java.vm.version"));
		System.out.println("Data model: " + Integer.getInteger("sun.arch.data.model"));
		
		long tenuredSize = Runtime.getRuntime().maxMemory();
		concurentMode = false;
		
		// getting more accurate data
		for(MemoryPoolMXBean bean : ManagementFactory.getMemoryPoolMXBeans()) {
			if (CONC_MODE.contains(bean.getName())) {
				concurentMode = true;
			}
			if (OLD_POOLS.contains(bean.getName())) {				
				tenuredSize = bean.getUsage().getMax();
				if (tenuredSize < 0) {
					tenuredSize = bean.getUsage().getCommitted();
				}
				System.out.println("Exact old space size is " + (tenuredSize) + " bytes");
				oldMemPool = bean;
				break;
			}
		}
		
		beans:
		for(GarbageCollectorMXBean bean :  ManagementFactory.getGarbageCollectorMXBeans()) {
			for(String pool: bean.getMemoryPoolNames()) {
				if (OLD_POOLS.contains(pool)) {
					oldGcMBean = (com.sun.management.GarbageCollectorMXBean) bean;
					continue beans;
				}
				else {
				}
			}
			youngGcMBean = bean;
		}

		boolean entrySizeAdjust = false;
		if (entrySize <= 0) {
			entrySizeAdjust = true;
			System.out.println("Estimating entry memory footprint ...");
			System.gc();
			long used = getOldSpaceUsed();
			
			int testSize = 100000;
			initMaps(testSize);
			
	 		while(size() < testSize) {
	 			populateMap(concurentMode, testSize);
			}
	 		System.gc();
	 		long footPrint = getOldSpaceUsed() - used;
	 		entrySize = align((int) (footPrint / testSize), 32);
	 		System.out.println("Entry footprint: " + entrySize);
	 		maps = null;
		}
		
		System.gc();
		long oldSpaceUsed = getOldSpaceUsed();
		long freeTenured = tenuredSize - getOldSpaceUsed();
		calculateCount(freeTenured);
		
		if (concurentMode) {
			System.out.println("Concurent mode is enabled");
			System.out.println("Available old space: " + (freeTenured >> 20) + "MiB");
		}
		else {
			System.out.println("Available old space: " + (freeTenured >> 20) + "MiB (-" + headRoom + " MiB)");
		}
		if (count < 0) {
			System.out.println("Heap size is too small, increase heap size or reduce headroom");
			return null;
		}
		System.out.println("Young space collector: " + youngGcMBean.getName());
		System.out.println("Old space collector: " + oldGcMBean.getName());
		
		System.out.println("Populating - " + count);
		initMaps(count);
		int n = 0;
		if (entrySizeAdjust) {
			int targetSize = 4 * count / 5;
			while(size() < targetSize) {
				populateMap(concurentMode, count);
				n++;
			}
			System.gc();
			System.gc();
			long footPrint = getOldSpaceUsed() - oldSpaceUsed;
			entrySize = align((int) (footPrint / size()), 32);
			System.out.println("Adjusted entry footprint: " + entrySize);
			calculateCount(freeTenured);
		}
		
 		while(size() < count) {
 			populateMap(concurentMode, count);
 			n++;
		}
 		
 		if (concurentMode) {
 			while(--n > 0) {
 				processMap(false);
 			}
 		}
 		
 		// Let's wait for at least one major collection to complete
 		if (!oldGcMBean.getName().startsWith("G1")) {
 			// in G1 incremental collections are not treated as full
 			// so we have to ignore this
	 		if (oldGcMBean != null) {
	 			long c = oldGcMBean.getCollectionCount();
	 			while(c == oldGcMBean.getCollectionCount()) {
	 				processMap(false);
	 			}
	 		}
 		}
 		
 		System.out.println("Size: " + size());
 		if (!dryMode) {
 			System.out.println("Processing ... ");
 		}
 		else {
 			System.out.println("Processing ... (DRY MODE ENABLED)");
 		}
 		StringBuffer sb = new StringBuffer();
 		sb.append("Limits:");
 		if (maxTime > 0) {
 			sb.append(" ").append(maxTime + " sec");
 		}
 		if (maxOld > 0) {
 			sb.append(" ").append(maxOld + " old collections");
 		}
 		if (maxYoung > 0) {
 			sb.append(" ").append(maxYoung + " young collections");
 		}
 		System.out.println(sb.toString());
 		
 		activeOverrideRate = overrideRate;
 		
 		YoungGcTimeTracker tracker = new YoungGcTimeTracker();
 		tracker.init();

 		// start count down here
 		long startTime = System.currentTimeMillis();
 		long oldC = oldGcMBean.getCollectionCount();
 		long youngC = youngGcMBean.getCollectionCount();
 		
 		while(true) {
			processMap(dryMode);
			tracker.probe();

			if (maxOld > 0) {
				if (oldGcMBean.getCollectionCount() - oldC > maxOld) {
					break;
				}
			}
			if (maxYoung > 0) {
				if (youngGcMBean.getCollectionCount() - youngC > maxYoung) {
					break;
				}
			}
			if (maxTime > 0 && (System.currentTimeMillis() > (startTime + TimeUnit.SECONDS.toMillis(maxTime)))) {
				break;
			}
		}
		
 		System.out.println("Benchmark complete");
		return tracker.result();
	}

	private long getOldSpaceUsed() {
		long usage = oldMemPool.getUsage().getUsed();
		if (usage < 0) {
			return oldGcMBean.getLastGcInfo().getMemoryUsageAfterGc().get(oldMemPool.getName()).getUsed();
		}
		else {
			return usage;
		}
		//return jstatLong("sun.gc.generation.1.space.0.used");
	}

	private void calculateCount(long tenuredSize) {
		count = (int) ((tenuredSize - (headRoom << 20)) / entrySize);
		if (concurentMode) {
			count /= 2;
		}
	}

	@SuppressWarnings("unchecked")
	private void initMaps(int entryCount) {
		maps = new Map[(entryCount + 200000 - 1) / 200000];
		for(int i = 0; i != maps.length; ++i) {
			maps[i] = new HashMap<Integer, Object>(250000 >> 8);
		}
	}

	private void processMap(boolean dry) {
		boolean remove = size() > 1.01 * count;
		for(int i = 0; i != 1000; ++i) {
			if ((remove) && random.nextBoolean()) {
				if (dry) {
					dryRemoveRandom(count);
				}
				else {
					removeRandom(count);
				}
			}
			else {
				if (dry) {
					dryPutRandom(count);
				}
				else {
					putRandom(count);
				}
			}
		}
	}

	private void populateMap(boolean concurentMode, int count) {
		for(int i = 0; i != 1000; ++i) {
			putRandom(count);
			if (concurentMode & random.nextInt(10) > 7) {
				removeRandom(count);
			}
		}
	}

	@SuppressWarnings("rawtypes")
	private int size() {
		int size = 0;
		for(Map map: maps) {
			size += map.size();
		}
		return size;
	}
	
	private Object newObject() {
		switch (mode) {
		case STRING: return new String(STRING_TEMPLATE);
		case INT: return new Integer(random.nextInt());
		case LONG: return new Long(random.nextInt());
		case CONST: return this;
		}
		return null;
	}
	
	private void putRandom(int count) {
		int key = random.nextInt(2 * count);
		if (Math.abs(random.nextDouble()) < activeOverrideRate) {
			Object val = newObject();		
			maps[key % maps.length].put(new Integer(key), val);
		}
		else {
			Integer ik = new Integer(key);
			Object v = maps[key % maps.length].get(ik);
			if (v != null) {
				maps[key % maps.length].put(ik, v);
			}
		}
	}

	private void dryPutRandom(int count) {
		int key = random.nextInt(2 * count);
		Object val = newObject();
		val.equals(maps[key % maps.length].get(key));
	}

	private void removeRandom(int count) {
		int key = random.nextInt(2 * count);
		maps[key % maps.length].remove(key);
	}

	private void dryRemoveRandom(int count) {
		int key = random.nextInt(2 * count);
		maps[key % maps.length].get(key);
	}

	private class YoungGcTimeTracker {

		private long totalTime = 0;
		private long evenCount = 0;
		private double squareTotal = 0;
		
		private long lastTime;
		private long lastYoungCount;
		private long lastOldCount;
		
		public void init() {
			while(true) {
				long ygc = youngGcMBean.getCollectionCount();
				long ogc = oldGcMBean.getCollectionCount();
				long yt = youngGcMBean.getCollectionTime(); 
				if (youngGcMBean.getCollectionCount() == ygc
						|| oldGcMBean.getCollectionCount() == ogc) {
					lastTime = yt;
					lastYoungCount = ygc;
					lastOldCount = ogc;
					break;
				}
			}
		}

		public void probe() {
			while(true) {
				long ygc = youngGcMBean.getCollectionCount();
				if (ygc == lastYoungCount) {
					return;
				}
				long ogc = oldGcMBean.getCollectionCount();
				long yt = youngGcMBean.getCollectionTime(); 
				if (youngGcMBean.getCollectionCount() == ygc
						|| oldGcMBean.getCollectionCount() == ogc) {
					
					long ycd = ygc - lastYoungCount;
					long ocd = ogc - lastOldCount;
					long td = yt - lastTime;
					
					lastYoungCount = ygc;
					lastOldCount = ogc;
					lastTime = yt;
					
					if (!concurentMode && ocd > 0) {
						// ignoring young part of full gc
						ycd -= ocd;						
					}
					
					if (ycd > 0) {
						totalTime += td;
						evenCount += ycd;
						double avg = ((double)td)/ycd;
						squareTotal += ycd * avg * avg;
						
						if (printEvents) {
							StringBuilder sb = new StringBuilder();
							sb.append("Young GC (" + ycd + " events), total time: " + td + "ms, (Old events: " + ogc + ")");
							System.out.println(sb);
						}
					}
					
					break;
				}
			}
		}
		
		public TestResult result() {
			TestResult result = new TestResult();
			result.totalSquareTime = squareTotal;
			result.totalTime = totalTime;
			result.youngGcCount = evenCount;
			return result;
		}
		
	}
	
	public static class TestResult implements Serializable {
		
		private static final long serialVersionUID = 20130518L;
		
		public long totalTime;
		public double totalSquareTime;
		public long youngGcCount;
		
		public double getAverage() {
			double avg = ((double)totalTime) / youngGcCount;
			return avg;
		}

		public double getStdDev() {
			double avg = ((double)totalTime) / youngGcCount;
			double stdDev = Math.sqrt((totalSquareTime / youngGcCount) - (avg * avg));
			return stdDev;
		}
		
		public String toString() {
			double avg = ((double)totalTime) / youngGcCount;
			double stdDev = Math.sqrt((totalSquareTime / youngGcCount) - (avg * avg));
			return String.format("%f [%f] ms", avg, stdDev); 
		}
	}
}
